from langchain.chat_models import ChatOpenAI
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.callbacks.manager import CallbackManager
from langchain_core.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

from config.model_config import get_chat_openai_zhipu

messages = [
    AIMessage(content="Hi."),
    SystemMessage(content="你的角色是一个诗人."),
    HumanMessage(content="用七言绝句的形式写一首关于AI的诗."),
]

streaming_chat = get_chat_openai_zhipu (
    streaming=True,
    callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]),
)

streaming_chat(messages)